import torch
from torch import nn
from torchsummary import summary

class VGG16(nn.Module):
    def __init__(self):
        super(VGG16, self).__init__()
        #定义一个块，Sequential是序列的意思。block1-block5提取特征
        self.block1=nn.Sequential(
            nn.Conv2d(1,64,3,1,1),#灰度图像
            nn.ReLU(),
            nn.Conv2d(64,64,3,1,padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2,2)
        )
        self.block2= nn.Sequential(
            nn.Conv2d(64, 128, 3, 1, 1),  # 灰度图像
            nn.ReLU(),
            nn.Conv2d(128, 128, 3, 1,padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.block3= nn.Sequential(
            nn.Conv2d(128, 256, 3, 1, 1),  # 灰度图像
            nn.ReLU(),
            nn.Conv2d(256, 256, 3, 1,padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 256, 3, 1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.block4 = nn.Sequential(
            nn.Conv2d(256, 512, 3, 1, 1),  # 灰度图像
            nn.ReLU(),
            nn.Conv2d(512, 512, 3, 1, padding=1),
            nn.ReLU(),
            nn.Conv2d(512, 512, 3, 1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.block5 = nn.Sequential(
            nn.Conv2d(512, 512, 3, 1, 1),  # 灰度图像
            nn.ReLU(),
            nn.Conv2d(512, 512, 3, 1, padding=1),
            nn.ReLU(),
            nn.Conv2d(512, 512, 3, 1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        #全连接层
        self.block6 = nn.Sequential(
           nn.Flatten(),
           nn.Linear(512 * 7 * 7, 4096),
           nn.ReLU(),
           nn.Linear(4096, 4096),
           nn.ReLU(),
           nn.Linear(4096, 10),
        )

        for m in self.modules():#给权重和偏置赋初值，一般pytorch会自动初始化
            #print(m)#打印出每层
            if isinstance(m, nn.Conv2d):#卷积层，特征提取
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')#
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)#b赋值0
            elif isinstance(m, nn.Linear):#全连接层
                nn.init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)#全连接层b

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.block4(x)
        x = self.block5(x)
        x = self.block6(x)
        return x

if __name__ == '__main__':
    # 这个是对的，不能device='cuda' if torch.cuda.is_available() else 'cpu'
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = VGG16().to(device)#实例化并且放到设备里
    print(summary(model, (1, 224, 224)))

